#!/usr/bin/env python

import sys
from PyQt4 import QtCore, QtGui

import simpletreemodel_rc

class TreeItem:
    def __init__(self, data, parent=None):
        self.parentItem = parent
        self.itemData = data
        self.childItems = []

    def appendChild(self, item):
        self.childItems.append(item)

    def child(self, row):
        return self.childItems[row]

    def childCount(self):
        return len(self.childItems)

    def columnCount(self):
        return len(self.itemData)

    def data(self, column):
        return self.itemData[column]

    def parent(self):
        return self.parentItem

    def row(self):
        if self.parentItem:
            return self.parentItem.childItems.index(self)
        return 0

class TreeModel(QtCore.QAbstractItemModel):
    def __init__(self,rootItem,parent=None):
        QtCore.QAbstractItemModel.__init__(self, parent)
        self.rootItem = rootItem

    def columnCount(self, parent):
        if parent.isValid():
            return parent.internalPointer().columnCount()
        else:
            return self.rootItem.columnCount()

    def data(self, index, role):
        if not index.isValid():
            return QtCore.QVariant()

        if role != QtCore.Qt.DisplayRole:
            return QtCore.QVariant()

        item = index.internalPointer()

        return QtCore.QVariant(item.data(index.column()))

    def flags(self, index):
        if not index.isValid():
            return QtCore.Qt.ItemIsEnabled

        return QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable

    def headerData(self, section, orientation, role):
        if orientation == QtCore.Qt.Horizontal and role == QtCore.Qt.DisplayRole:
            return self.rootItem.data(section)

        return QtCore.QVariant()

    def index(self, row, column, parent):
        if row < 0 or column < 0 or row >= self.rowCount(parent) or column >= self.columnCount(parent):
            return QtCore.QModelIndex()

        if not parent.isValid():
            parentItem = self.rootItem
        else:
            parentItem = parent.internalPointer()

        childItem = parentItem.child(row)
        if childItem:
            return self.createIndex(row, column, childItem)
        else:
            return QtCore.QModelIndex()

    def parent(self, index):
        if not index.isValid():
            return QtCore.QModelIndex()

        childItem = index.internalPointer()
        parentItem = childItem.parent()

        if parentItem == self.rootItem:
            return QtCore.QModelIndex()

        return self.createIndex(parentItem.row(), 0, parentItem)

    def rowCount(self, parent):
        if parent.column() > 0:
            return 0

        if not parent.isValid():
            parentItem = self.rootItem
        else:
            parentItem = parent.internalPointer()

        return parentItem.childCount()

def testData():
    rootData = []
    rootData.append(QtCore.QVariant("Title"))
    rootData.append(QtCore.QVariant("Summary"))
    rootData = TreeItem(rootData)

    for data in range(12):
        child = TreeItem((str(data),str(data)),rootData)
        rootData.appendChild(child)
    return rootData

################################################################

class Entity():
    def __init__(self,name,desc,parent=None):
        self.name=name
        self.desc=desc
        self.parent=parent

def create_entity_list():
    e1 = Entity("work","working")
    e2 = Entity("study","studying")
    e3 = Entity("pts","ppp",e1)
    e4 = Entity("spi","sss",e3)
    return set((e1,e2,e3,e4))

def testData2():
    rootData = TreeItem( ('name','desc') ,None)
    elist = create_entity_list()
    return rootData

if __name__ == "__main__":
    app = QtGui.QApplication(sys.argv)

    model = TreeModel(testData())

    view = QtGui.QTreeView()
    view.setModel(model)
    view.setWindowTitle("Simple Tree Model")
    view.show()
    sys.exit(app.exec_())
